{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "903c3461-0c63-4790-9e8c-75aeee52330d",
   "metadata": {},
   "source": [
    "### Advanced RAG Testbed for ORAN chatbot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba1aae2f-43ba-42da-9ac7-ddf68075aec9",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# !pip install xlrd\n",
    "# !pip install werkzeug\n",
    "# !pip install markdown\n",
    "# !pip install chromadb\n",
    "# !pip install helper-utils\n",
    "# !pip install langchain\n",
    "# !pip install chromadb\n",
    "# !pip install openai\n",
    "# !pip install sentence_transformers\n",
    "# !pip install pandas\n",
    "# !pip install pypdf\n",
    "# !pip install openpyxl\n",
    "# !pip install ragas --upgrade\n",
    "\n",
    "# !pip install -r requirements.txt\n",
    "# # !pip install --upgrade langchain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbc785ce-0418-4913-8a4a-464b388e05b1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from langchain.text_splitter import RecursiveCharacterTextSplitter, SentenceTransformersTokenTextSplitter\n",
    "from langchain_community.document_loaders import DirectoryLoader\n",
    "from langchain.document_loaders import TextLoader, PyPDFLoader\n",
    "from langchain_community.document_loaders import UnstructuredHTMLLoader\n",
    "from langchain_community.document_loaders import UnstructuredPDFLoader\n",
    "from langchain_community.document_loaders import PyPDFLoader\n",
    "from langchain.document_loaders import UnstructuredFileLoader, Docx2txtLoader\n",
    "from langchain.vectorstores import FAISS\n",
    "import pickle\n",
    "from langchain.embeddings import HuggingFaceEmbeddings\n",
    "import torch.nn.functional as F\n",
    "import torch\n",
    "from torch import Tensor\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "from chromadb import Documents, EmbeddingFunction, Embeddings\n",
    "import chromadb\n",
    "import openai\n",
    "from openai._client import OpenAI\n",
    "from dotenv import load_dotenv, find_dotenv\n",
    "import numpy as np\n",
    "from sentence_transformers import CrossEncoder\n",
    "import re\n",
    "import pandas as pd\n",
    "import json\n",
    "from IPython.display import JSON\n",
    "from typing import List\n",
    "\n",
    "from langchain.chains import LLMChain\n",
    "from langchain.output_parsers import PydanticOutputParser\n",
    "from langchain.prompts import PromptTemplate\n",
    "from pydantic import BaseModel, Field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a33bf3ef-aaf3-4322-8824-db90bb6d2bce",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "# from triton_trt_llm import TensorRTLLM\n",
    "# trtllm =TensorRTLLM(server_url=\"llm:8001\", model_name=\"ensemble\", tokens=300)\n",
    "\n",
    "DOCS_DIR = \"../ORAN_kb/\" #\n",
    "DOCS_DIR = \"../ORAN_kb/\"\n",
    "from transformers import AutoTokenizer, AutoModel\n",
    "############################################\n",
    "# Component #2 - Embedding Model and LLM\n",
    "############################################\n",
    "os.environ['NVIDIA_API_KEY'] = \"\"\n",
    "os.environ['NVAPI_KEY'] = \"\"\n",
    "from langchain_nvidia_ai_endpoints import ChatNVIDIA, NVIDIAEmbeddings"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa4a26a1-8fa5-4c53-9820-318f120d6f6d",
   "metadata": {},
   "source": [
    "### Load Documents\n",
    "Follow the step number 1 [defined here](../notebooks/05_dataloader.ipynb) to upload the pdf's to Milvus server."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61eac811-00f6-425f-8c09-276b65d9cf4d",
   "metadata": {},
   "outputs": [],
   "source": [
    "DOCS_DIR = \"../ORAN_kb\"\n",
    "# Load raw documents from the directory\n",
    "text_loader_kwargs={'autodetect_encoding': True} #loader_kwargs=text_loader_kwargs\n",
    "raw_txts = DirectoryLoader(DOCS_DIR, glob=\"**/*.txt\", show_progress=True, loader_cls=TextLoader).load()\n",
    "raw_htmls = DirectoryLoader(DOCS_DIR, glob=\"**/*.html\", show_progress=True, loader_cls=UnstructuredHTMLLoader).load()\n",
    "raw_pdfs = DirectoryLoader(DOCS_DIR, glob=\"**/*.pdf\", show_progress=True, loader_cls=UnstructuredPDFLoader).load()\n",
    "raw_docs = DirectoryLoader(DOCS_DIR, glob=\"**/*.docx\", show_progress=True, loader_cls=Docx2txtLoader).load()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f604224-3cf2-46c7-bf6a-f2083b14aecc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def remove_line_break(text):\n",
    "    text = text.replace(\"\\n\", \" \").strip()\n",
    "    text = re.sub(\"\\.\\.+\", \"\", text)\n",
    "    text = re.sub(\" +\", \" \", text)\n",
    "    return text\n",
    "    \n",
    "def remove_two_points(text):\n",
    "    text = text.replace(\"..\",\"\")\n",
    "    return text\n",
    "    \n",
    "def remove_two_slashes(text):\n",
    "    text = text.replace(\"__\",\"\")\n",
    "    return text\n",
    "    \n",
    "def just_letters(text):\n",
    "    return re.sub(r\"[^a-z]+\", \"\", text).strip()\n",
    "    \n",
    "def remove_non_english_letters(text):\n",
    "    return re.sub(r\"[^\\x00-\\x7F]+\", \"\", text)\n",
    "    \n",
    "def langchain_length_function(text):\n",
    "    return len(just_letters(remove_line_break(text)))\n",
    "    \n",
    "def word_count(text):\n",
    "    text = remove_line_break(text)\n",
    "    text = just_letters(text)\n",
    "    tokenizer = RegexpTokenizer(r\"[A-Za-z]\\w+\")\n",
    "    tokenized_text = tokenizer.tokenize(text)\n",
    "    tokenized_text_len = len(tokenized_text)\n",
    "    return tokenized_text_len\n",
    "    \n",
    "def truncate(text, max_word_count=1530):\n",
    "    tokenizer = RegexpTokenizer(r\"[A-Za-z]\\w+\")\n",
    "    tokenized_text = tokenizer.tokenize(text)\n",
    "    return \" \".join(tokenized_text[:max_word_count])\n",
    "\n",
    "def strip_non_ascii(string):\n",
    "    ''' Returns the string without non ASCII characters'''\n",
    "    stripped = (c for c in string if 0 < ord(c) < 127)\n",
    "    return ''.join(stripped)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b90c7e6b-9e54-4727-a331-2fb7f45be1fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "text_splitter = RecursiveCharacterTextSplitter(separators=[\"\\n\\n\", \"\\n\", \". \", \" \", \"\"], chunk_size=500, chunk_overlap=100, length_function=langchain_length_function, )\n",
    "all_documents=raw_pdfs+raw_docs\n",
    "\n",
    "# split all the docs\n",
    "docs = text_splitter.split_documents(all_documents)\n",
    "documents = docs\n",
    "#remove short chuncks\n",
    "filtered_documents = [item for item in documents if len(item.page_content) >= 200]\n",
    "documents = filtered_documents\n",
    "#remove line break\n",
    "for i in range(0,len(documents)-1):\n",
    "    documents[i].page_content=remove_line_break(documents[i].page_content)\n",
    "#remove two points\n",
    "for i in range(0,len(documents)-1):\n",
    "    documents[i].page_content=remove_two_points(documents[i].page_content)\n",
    "#remove non english characters points\n",
    "for i in range(0,len(documents)-1):\n",
    "    documents[i].page_content=remove_two_slashes(documents[i].page_content)\n",
    "#remove two points\n",
    "for i in range(0,len(documents)-1):\n",
    "    documents[i].page_content=remove_two_points(documents[i].page_content)\n",
    "\n",
    "orig_docs = docs\n",
    "print(len(docs))\n",
    "print(docs[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "844a8dea-9c83-49f8-bad9-0b5a90a50fbf",
   "metadata": {},
   "outputs": [],
   "source": [
    "metadatas = []\n",
    "character_split_texts = []\n",
    "prev_len = len(character_split_texts)\n",
    "for doc in docs:\n",
    "    # split text in each page\n",
    "    character_split_texts.append(doc.page_content)\n",
    "    \n",
    "    new_len = len(character_split_texts)\n",
    "    z = new_len - prev_len\n",
    "    #Store metadata associated with each chunk\n",
    "    for i in range(z):\n",
    "        metadatas.append(doc.metadata)\n",
    "    prev_len = new_len\n",
    "print(character_split_texts[100])\n",
    "print(metadatas[100])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84a96afc-a769-4b1e-8bda-450d4e6f85db",
   "metadata": {},
   "source": [
    "#### Generate and Store Embeddings\n",
    "##### a) Generate Embeddings \n",
    "[Embeddings](https://python.langchain.com/docs/modules/data_connection/text_embedding/) for documents are created by vectorizing the document text; this vectorization captures the semantic meaning of the text. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd2f11b0-6b74-4f2c-ae7a-c8c2635b5b1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Pooling function needed for embeddings\n",
    "def average_pool(last_hidden_states: Tensor,\n",
    "                 attention_mask: Tensor) -> Tensor:\n",
    "    last_hidden = last_hidden_states.masked_fill(~attention_mask[..., None].bool(), 0.0)\n",
    "    return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]\n",
    "\n",
    "#automatically get the tokenizer from the model folder\n",
    "tokenizer = AutoTokenizer.from_pretrained(\"../ft_paq_squad_nq_nli_nvolve40k_hn/hf_model\")\n",
    "#automatically get the model from the model folder\n",
    "model = AutoModel.from_pretrained(\"../ft_paq_squad_nq_nli_nvolve40k_hn/hf_model\").cuda()\n",
    "\n",
    "#Create custom embedding class for ChromaDB so that it works with NVIDIA NVOLVE40K embedding model\n",
    "class custom_embedding_NVIDIA(EmbeddingFunction[Documents]):\n",
    "    def __init__(self):\n",
    "        model = AutoModel.from_pretrained(\"../ft_paq_squad_nq_nli_nvolve40k_hn/hf_model\") # ---> HERE Load your custom model\n",
    "        tokenizer = AutoTokenizer.from_pretrained(\"../ft_paq_squad_nq_nli_nvolve40k_hn/hf_model\")\n",
    "        # self.model.cuda()\n",
    "    def embed_documents(self, input: Documents) -> Embeddings:\n",
    "        \n",
    "        batch_dict = tokenizer(input, max_length=512, padding=True, truncation=True, return_tensors='pt')\n",
    "        with torch.no_grad():\n",
    "            outputs = model(**batch_dict)\n",
    "        embeddings = average_pool(outputs.last_hidden_state, batch_dict['attention_mask'])\n",
    "\n",
    "        # normalize embeddings\n",
    "        embeddings = F.normalize(embeddings, p=2, dim=1)\n",
    "        embeddings = [embed.detach().numpy().tolist() for embed in embeddings]\n",
    "        return embeddings\n",
    "\n",
    "# custom = custom_embedding_NVIDIA()\n",
    "nv_document_embedder = NVIDIAEmbeddings(model=\"nvolveqa_40k\", model_type=\"passage\")\n",
    "\n",
    "model_name = \"../ft_paq_squad_nq_nli_nvolve40k_hn/hf_model\"\n",
    "model_kwargs = {\"device\": \"cuda\"}\n",
    "encode_kwargs = {\"normalize_embeddings\": True}\n",
    "nv_local_embeddings = HuggingFaceEmbeddings(\n",
    "    model_name=model_name,\n",
    "    model_kwargs=model_kwargs,\n",
    "    encode_kwargs=encode_kwargs\n",
    ")\n",
    "\n",
    "bge_embedder = HuggingFaceEmbeddings(\n",
    "    model_name=\"BAAI/bge-large-en-v1.5\",\n",
    "    model_kwargs=model_kwargs,\n",
    "    encode_kwargs=encode_kwargs\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "904cf7f1-d222-4158-9e3b-feffccb3a68e",
   "metadata": {},
   "outputs": [],
   "source": [
    "docs = orig_docs\n",
    "#Create vectordatabase\n",
    "vectorstore = FAISS.from_documents(docs, nv_local_embeddings)\n",
    "vector_store_path = \"../evals/vectorstore_nv_hf.pkl\"\n",
    "vector_store_exists = os.path.exists(vector_store_path)\n",
    "with open(vector_store_path, \"wb\") as f:\n",
    "    pickle.dump(vectorstore, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3b4694b-8928-4a80-b4a4-a6b658e54e4e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ef588e2-df44-4fe3-b9a5-a6e2e498f622",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Create vectordatabase\n",
    "vectorstore = FAISS.from_documents(docs, bge_embedder)\n",
    "vector_store_path = \"../vectorstore/oran/vectorstore_bge.pkl\"\n",
    "vector_store_exists = os.path.exists(vector_store_path)\n",
    "with open(vector_store_path, \"wb\") as f:\n",
    "    pickle.dump(vectorstore, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "deac9f07-d087-40c5-bfd6-90b36aa8e06b",
   "metadata": {},
   "outputs": [],
   "source": [
    "file = open(\"../evals/vectorstore_nv_hf.pkl\",'rb')\n",
    "# file = open(\"/home/jvamaraju/oran-chatbot/vectorstore/oran/vectorstore_bge.pkl\",'rb')\n",
    "vectorstore = pickle.load(file)\n",
    "file.close()\n",
    "print(len(vectorstore.docstore._dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59eb0e6a-f0fd-4331-b9f8-c60d5b6eae23",
   "metadata": {},
   "outputs": [],
   "source": [
    "os.environ['NVIDIA_API_KEY'] = \"\"\n",
    "os.environ['NVAPI_KEY'] = \"\"\n",
    "os.environ['NGC_API_KEY'] = \"\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a8a2796-b6af-4b37-ae29-aec95dd90372",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from langchain.retrievers.multi_query import MultiQueryRetriever\n",
    "prompt_template = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"You are a helpful AI assistant named Envie. You will reply to questions only based on the context that you are provided. If something is out of context, you will refrain from replying and politely decline to respond to the user.\"), (\"user\", \"{input}\")]\n",
    ")\n",
    "user_input = \"What are the beamforming methods specified in O-RAN?\"\n",
    "retriever = vectorstore.as_retriever(search_type=\"similarity_score_threshold\", search_kwargs={\"k\": 6, \"score_threshold\": 0.3})\n",
    "\n",
    "\n",
    "docs = retriever.get_relevant_documents(user_input)\n",
    "context = \"\"\n",
    "for doc in docs:\n",
    "    context += doc.page_content + \"\\n\\n\"\n",
    "\n",
    "# print(context) \n",
    "\n",
    "def nemo_rag(query, retrieved_documents, model=\"playground_llama2_70b\"):\n",
    "    #Combine the query and retrieved documents and send to model\n",
    "    llm = ChatNVIDIA(model=model)\n",
    "    prompt_template = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"You are a helpful expert ORAN assistant. Your task is to provide a precise and accurate response that directly addresses user's question using the given relevant information. First, meticulously analyze and understand the provided relevant information in detail. With a thorough understanding of both the question and the relevant information, provide a response that is clear, detailed, and accurate that directly addresses user's query. Ensure that your response is strictly based on the provided relevant information. Your response is crucial for my career; hence, accuracy is of utmost importance. So, take a deep breath and work on this task step-by-step.\"), (\"user\", \"{input}\")]\n",
    "        )\n",
    "    chain = prompt_template | llm | StrOutputParser()\n",
    "    augmented_user_input = \"Context: \" + retrieved_documents + \"\\n\\nQuestion: \" + query + \"\\n\"\n",
    "    full_response = \"\"\n",
    "    for response in chain.stream({\"input\": augmented_user_input}):\n",
    "        full_response += response\n",
    "    final_ans = full_response\n",
    "    return final_ans\n",
    "\n",
    "print(nemo_rag(user_input, context))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c31cfcc5-3acf-487c-a30b-b7486fa929d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "user_input = \"What are the beamforming methods?\"\n",
    "# user_input = \"What are the compression methods?\"\n",
    "# user_input = \"What are the NB-IOT downlink use cases addressed by WG4?\"\n",
    "# user_input = \"What are the LLS configurations for S-plane?\"\n",
    "# Output parser will split the LLM result into a list of queries\n",
    "class LineList(BaseModel):\n",
    "    # \"lines\" is the key (attribute name) of the parsed output\n",
    "    lines: List[str] = Field(description=\"Lines of text\")\n",
    "\n",
    "\n",
    "class LineListOutputParser(PydanticOutputParser):\n",
    "    def __init__(self) -> None:\n",
    "        super().__init__(pydantic_object=LineList)\n",
    "\n",
    "    def parse(self, text: str) -> LineList:\n",
    "        lines = text.strip().split(\"\\n\")\n",
    "        return LineList(lines=lines)\n",
    "\n",
    "\n",
    "output_parser = LineListOutputParser()\n",
    "\n",
    "QUERY_PROMPT = PromptTemplate(\n",
    "    input_variables=[\"question\"],\n",
    "    template=\"\"\"You are an ORAN standards assistant. Your task is to generate five \n",
    "    different versions of the given user question relevant to ORAN from ORAN documents. By generating multiple perspectives on the user question, your goal is to help\n",
    "    the user overcome some of the limitations of the distance-based similarity search. \n",
    "    Provide these alternative questions separated by newlines.\n",
    "    Original question: {question}\"\"\",\n",
    ")\n",
    "llm = ChatNVIDIA(model=\"playground_llama2_70b\")\n",
    "\n",
    "# Chain\n",
    "llm_chain = LLMChain(llm=llm, prompt=QUERY_PROMPT, output_parser=output_parser)\n",
    "\n",
    "# Other inputs\n",
    "retriever_mq = MultiQueryRetriever(\n",
    "    retriever=vectorstore.as_retriever(search_type=\"similarity_score_threshold\",search_kwargs={\"k\": 10, \"score_threshold\": 0.2}), llm_chain=llm_chain, parser_key=\"lines\"\n",
    ") \n",
    "unique_docs = retriever_mq.get_relevant_documents(query=user_input)\n",
    "#Remove all duplicated documents and retain the original metadata\n",
    "unique_documents = []\n",
    "unique_documents_metadata = []\n",
    "for document in unique_docs:\n",
    "    if document not in unique_documents:\n",
    "        unique_documents.append(document.page_content)\n",
    "        unique_documents_metadata.append(document.metadata['source'])\n",
    "print(\"length of unique docs: \", len(unique_documents))\n",
    "cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2') # ('BAAI/bge-reranker-large')('cross-encoder/ms-marco-MiniLM-L-6-v2')\n",
    "pairs = [[user_input, doc] for doc in unique_documents]\n",
    "scores = cross_encoder.predict(pairs)\n",
    "#Sort the scores from highest to least\n",
    "order_ids =  np.argsort(scores)[::-1]\n",
    "# print(order_ids)\n",
    "new_updated_documents = []\n",
    "new_updated_sources = []\n",
    "#Get the top 6 scores\n",
    "if len(order_ids)>=10:\n",
    "    for i in range(10):\n",
    "        new_updated_documents.append(unique_documents[order_ids[i]])\n",
    "        new_updated_sources.append(unique_documents_metadata[order_ids[i]])\n",
    "else:\n",
    "    for i in range(len(order_ids)):\n",
    "        new_updated_documents.append(unique_documents[order_ids[i]])\n",
    "        new_updated_sources.append(unique_documents_metadata[order_ids[i]])\n",
    "\n",
    "print(\"length of unique docs: \", len(new_updated_documents))\n",
    "context = \"\"\n",
    "for doc in new_updated_documents:\n",
    "    context += doc+ \" \"\n",
    "# print(context)\n",
    "print(nemo_rag(user_input, context))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe4814ad-badf-4961-85d4-3f48ce0e4efa",
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_multiple_query(query, model=\"playground_llama2_70b\"):\n",
    "    #For the given query, lets create 5 additional queries using the LLM\n",
    "    llm = ChatNVIDIA(model=model,max_output_token=500, top_k=1, top_p=0.0)\n",
    "    prompt_template = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"You are an expert in the field of Oran specifications and processes. User has a question related to this field, sourced from relevant documents.\\nTo help the user find the information they need, please suggest five additional related questions. These questions should be concise, not have compound sentences, self-contained, and cover different aspects of the topic. Each question should be complete and relevant to the original query.\\nPlease output one question per line without numbering them.\"), (\"user\", \"{input}\")]\n",
    "        )\n",
    "    chain = prompt_template | llm | StrOutputParser()\n",
    "    augmented_user_input = \"\\n\\nQuestion: \" + query + \"\\n\"\n",
    "    full_response = \"\"\n",
    "    for response in chain.stream({\"input\": augmented_user_input}):\n",
    "        full_response += response\n",
    "    final_ans = full_response\n",
    "    final_ans = final_ans.split(\"\\n\")\n",
    "    return final_ans\n",
    "    \n",
    "print(augment_multiple_query(user_input))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89f0a2f6-4cc6-4fa3-bf66-37682189682d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def augment_query_generated(query, model=\"playground_llama2_70b\"):\n",
    "    #For the given query, lets create a hypothetical answer using the LLM\n",
    "    llm = ChatNVIDIA(model=model,max_output_token=500, top_k=1, top_p=0.0)\n",
    "    prompt_template = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"You are an expert in the field of ORAN specifications and processes. Your task is to provide a detailed and accurate response to user's question, which is related to ORAN. Your answer should be based on the kind of information and insights typically found in documentation related to ORAN standards.\"), (\"user\", \"{input}\")]\n",
    "        )\n",
    "    chain = prompt_template | llm | StrOutputParser()\n",
    "    augmented_user_input = \"\\n\\nQuestion: \" + query + \"\\n\"\n",
    "    full_response = \"\"\n",
    "    for response in chain.stream({\"input\": augmented_user_input}):\n",
    "        full_response += response\n",
    "    final_ans = full_response\n",
    "    return final_ans\n",
    "\n",
    "print(augment_query_generated(user_input))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "538e385e-5629-4ca1-a97a-ff94971c70d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "def query_rewriting(query, history, model=\"playground_llama2_70b\"):\n",
    "    #Rewrite the given query using the context from LLM\n",
    "    llm = ChatNVIDIA(model=model)\n",
    "    prompt_template = ChatPromptTemplate.from_messages(\n",
    "    [(\"system\", \"Here is the conversation history between user and Assistant. You are an expert in the field of ORAN standards and specifications. User has a follow-up question to the conversation. Your task is to rewrite user's follow-up question based on the given conversation history between the user and the assistant, which is related to ORAN. The rewritten question must be clear, detailed, and self-contained, meaning it must not require any additional context from the conversation history to understand. Ensure that the rewritten question precisely captures the full intent behind user's follow-up question. Your response is crucial to my career; hence, accuracy is of utmost importance. So, take a deep breath and work on this task step-by-step.\"), (\"user\", \"{input}\")]\n",
    "        )\n",
    "    chain = prompt_template | llm | StrOutputParser()\n",
    "    augmented_user_input = \"History: \" + history + \"\\n\\nQuestion: \" + query + \"\\n\"\n",
    "    full_response = \"\"\n",
    "    for response in chain.stream({\"input\": augmented_user_input}):\n",
    "        full_response += response\n",
    "    final_ans = full_response\n",
    "    return final_ans"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3a60ccd-2942-4b11-9d14-28d5da62f749",
   "metadata": {},
   "source": [
    "### lets try base RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68d57f11-2778-45e7-aed1-01226dab085a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def base_rag_workflow(query):\n",
    "    #Get top 6 retrieved chunks for the query\n",
    "    retriever = vectorstore.as_retriever(search_kwargs={\"k\": 6, \"score_threshold\": 0.5})\n",
    "    docs = retriever.get_relevant_documents(query)\n",
    "    context = \"\"\n",
    "    retrieved_documents = []\n",
    "    retrieved_metadata = []\n",
    "    for doc in docs:\n",
    "        context += doc.page_content + \"\\n\\n\"\n",
    "        retrieved_documents.append(doc.page_content)\n",
    "        retrieved_metadata.append(doc.metadata)\n",
    "    #Get all the results and metadatas associated with each result\n",
    "    #Send the top 6 results along with query to LLM\n",
    "    output = nemo_rag(query=query, retrieved_documents=context)\n",
    "    return output,retrieved_documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1e1a7fd-65b3-4284-9b24-8bcb6a007570",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# load the evaluation data\n",
    "f = open('syn_data_oran.json')\n",
    "syn_data = json.load(f)\n",
    "for count,entry in enumerate(syn_data):\n",
    "    query = entry['question']\n",
    "    ans, cons = base_rag_workflow(query)\n",
    "    entry[\"answer\"] = ans\n",
    "    entry[\"contexts\"] = cons\n",
    "    print(count, query)\n",
    "\n",
    "with open('eval_syn_QA_NV_hf.json', 'w') as f:\n",
    "    json.dump(syn_data, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74d59fbb-e32b-4272-b09e-642edb281569",
   "metadata": {},
   "source": [
    "### lets try multiple query based RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65e9c49c-ce0f-4291-a522-a7a867a12a42",
   "metadata": {},
   "outputs": [],
   "source": [
    "def augmented_queries_answer(query):\n",
    "    #Create multiple queries\n",
    "    augmented_queries = augment_multiple_query(query)\n",
    "    #Create a list of queries\n",
    "    queries = [query] + augmented_queries[2:]\n",
    "    print(\"Queries are = \", queries)\n",
    "    # #Get all the retrievals for each queries\n",
    "    # #Get all the results and metadatas associated with each result\n",
    "    #Get top 10 retrieved chunks for the queries\n",
    "    retriever = vectorstore.as_retriever(search_type=\"similarity_score_threshold\",search_kwargs={\"k\": 6, \"score_threshold\": 0.3})\n",
    "    retrieved_documents = []\n",
    "    retrieved_metadatas = []\n",
    "\n",
    "    for query in queries:\n",
    "        docs = retriever.get_relevant_documents(query)\n",
    "        for doc in docs:\n",
    "            retrieved_documents.append(doc.page_content)\n",
    "            retrieved_metadatas.append(doc.metadata)\n",
    "    print(\"length of retrieved docs: \", len(retrieved_documents))\n",
    "    #Remove all duplicated documents and retain the original metadata\n",
    "    unique_documents = []\n",
    "    unique_documents_metadata = []\n",
    "    for document,source in zip(retrieved_documents,retrieved_metadatas):\n",
    "            if document not in unique_documents:\n",
    "                unique_documents.append(document)\n",
    "                unique_documents_metadata.append(source)\n",
    "    # unique_documents = retrieved_documents\n",
    "    # unique_documents_metadata = retrieved_metadatas\n",
    "    print(\"length of unique docs: \", len(unique_documents))\n",
    "    #Instantiate the cross-encoder model and get scores for each retrieved document\n",
    "    cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')\n",
    "    pairs = [[query, doc] for doc in unique_documents]\n",
    "    scores = cross_encoder.predict(pairs)\n",
    "    #Sort the scores from highest to least\n",
    "    order_ids =  np.argsort(scores)[::-1]\n",
    "    print(order_ids)\n",
    "    new_updated_documents = []\n",
    "    #Get the top 6 scores\n",
    "    if len(order_ids)>=6:\n",
    "        for i in range(6):\n",
    "            new_updated_documents.append(unique_documents[order_ids[i]])\n",
    "    else:\n",
    "        for i in range(len(order_ids)):\n",
    "            new_updated_documents.append(unique_documents[order_ids[i]])\n",
    "        \n",
    "    context = \"\"\n",
    "    for doc in new_updated_documents:\n",
    "            context += doc + \"\\n\\n\"\n",
    "    #Send the top 6 results along with the query to LLM\n",
    "    output = nemo_rag(query=query, retrieved_documents=context)\n",
    "    return output, new_updated_documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67dffe97-03ff-4c05-92e0-4adbce76e135",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "f = open('syn_data_oran.json')\n",
    "syn_data = json.load(f)\n",
    "for count,entry in enumerate(syn_data):\n",
    "    # history = \"\"\n",
    "    init_query = entry['question']\n",
    "    print(count, init_query)\n",
    "    ans, cons =  augmented_queries_answer(init_query)\n",
    "    # history += f\"query: {query}\\nAssistant: {ans}\\n\\n\" \n",
    "    entry[\"answer\"] = ans\n",
    "    entry[\"contexts\"] = cons\n",
    "\n",
    "with open('eval_syn_QA_NV_hf_multiplequery.json', 'w') as f:\n",
    "    json.dump(syn_data, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "484b2748-eb8a-42bc-9d16-640d90fa3e3f",
   "metadata": {},
   "source": [
    "### lets try HYDE (hypothetical llm answer) based RAG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a844d5fa-74da-4682-b7a8-3c53e76f7b66",
   "metadata": {},
   "outputs": [],
   "source": [
    "def augmented_hypothetical_answer(query):\n",
    "    #Create multiple queries\n",
    "    sample_response = augment_query_generated(query)\n",
    "    augmented_queries = query + \" \" + sample_response\n",
    "    print(\"Augmented query = \", augmented_queries)\n",
    "    #Get all the retrievals for each queries\n",
    "    #Get all the results and metadatas associated with each result\n",
    "    retriever = vectorstore.as_retriever(search_type=\"similarity_score_threshold\",search_kwargs={\"k\": 10, \"score_threshold\": 0.3})\n",
    "    retrieved_documents = []\n",
    "    retrieved_metadatas = []\n",
    "\n",
    "    docs = retriever.get_relevant_documents(augmented_queries)\n",
    "    for doc in docs:\n",
    "        retrieved_documents.append(doc.page_content)\n",
    "        retrieved_metadatas.append(doc.metadata)\n",
    "    print(\"length of retrieved docs: \", len(retrieved_documents))\n",
    "    #Remove all duplicated documents\n",
    "    unique_documents = []\n",
    "    unique_documents_metadata = []\n",
    "    for documents,sources in zip(retrieved_documents,retrieved_metadatas):\n",
    "        for document,source in zip(documents,sources):\n",
    "            if document not in unique_documents:\n",
    "                unique_documents.append(document)\n",
    "                unique_documents_metadata.append(source)\n",
    "\n",
    "    #Instantiate the cross-encoder model and get scores for each retrieved document\n",
    "    cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')\n",
    "    pairs = [[query, doc] for doc in unique_documents]\n",
    "    scores = cross_encoder.predict(pairs)\n",
    "    #Sort the scores from highest to lowest\n",
    "    order_ids =  np.argsort(scores)[::-1]\n",
    "    \n",
    "    new_updated_documents = []\n",
    "    #Get the top 6 scores\n",
    "    if len(order_ids)>=6:\n",
    "        for i in range(6):\n",
    "            new_updated_documents.append(unique_documents[order_ids[i]])\n",
    "    else:\n",
    "        for i in range(len(order_ids)):\n",
    "            new_updated_documents.append(unique_documents[order_ids[i]])\n",
    "    \n",
    "    context = \"\"\n",
    "    for doc in new_updated_documents:\n",
    "            context += doc + \"\\n\\n\"\n",
    "    \n",
    "    #Send the top 6 results along with query to LLM\n",
    "    output = nemo_rag(query=query, retrieved_documents=context)\n",
    "    return output, new_updated_documents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2b2c1ae-ddc8-49eb-8ada-f832b51b0a5c",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "f = open('syn_data_oran.json')\n",
    "syn_data = json.load(f)\n",
    "for count,entry in enumerate(syn_data):\n",
    "    # history = \"\"\n",
    "    init_query = entry['question']\n",
    "    print(count, init_query)\n",
    "    ans, cons =  augmented_hypothetical_answer(init_query)\n",
    "    # history += f\"query: {query}\\nAssistant: {ans}\\n\\n\" \n",
    "    entry[\"answer\"] = ans\n",
    "    entry[\"contexts\"] = cons\n",
    "\n",
    "with open('eval_syn_QA_NV_hf_hyde.json', 'w') as f:\n",
    "    json.dump(syn_data, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4f7a171-18e5-4b2c-8b22-d130baf50826",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
